import uuid
from typing import Dict, Set
from fastapi import WebSocket

from api.utils import app

clients: Dict[str, WebSocket] = {}

connected_clients: Dict[str, WebSocket] = {}

# 用于存储所有客户端ID的集合
client_ids: Set[str] = set()


@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    # 生成一个唯一的客户端ID
    client_id = str(uuid.uuid4())

    # 存储客户端ID和WebSocket实例的映射
    connected_clients[client_id] = websocket
    await websocket.accept()
    obj_dict = {
        "message": client_id,
        "type": 'clientid'
    }
    await websocket.send_json(obj_dict)
    try:
        while True:
            # 等待接收客户端的消息
            data = await websocket.receive_text()
            # 处理消息;
            # 广播消息到所有连接的客户端，除了发送者
            # for cid in connected_clients.keys():
            #     if cid != client_id:
            #         client = connected_clients[cid]
            #         await client.send_text(data)
    except Exception as e:
        print(f"WebSocket连接中断: {e}")
    finally:
        # 客户端断开连接时的处理
        connected_clients.pop(client_id, None)
        client_ids.remove(client_id)


async def send_message_to_client(client_id: str, message: str):
    if client_id in connected_clients:
        # 向特定客户端发送消息
        websocket = connected_clients[client_id]
        await websocket.send_text(message)
    else:
        print(f"No client with id {client_id} found.")
